import numpy as np
from sklearn.manifold import TSNE
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time
import json
import pdb

meta_type = 1

'''
shots = 10
shots_2 = 50
'''

for epoch in range(21, 22):
    '''
    file1 = 'model_path/pooled_feat/att_shots_%.0f_epoch_%.0f.json' % (shots, epoch)
    att_data1 = json.load(open(file1, 'r'))
    att_data1 = np.array(att_data1['class_attentions']).reshape(-1, 2048)
    file2 = 'model_path/pooled_feat/att_shots_%.0f_epoch_%.0f.json' % (shots_2, epoch)
    att_data2 = json.load(open(file2, 'r'))
    att_data2 = np.array(att_data2['class_attentions']).reshape(-1, 2048)
    att_data = np.vstack((att_data1, att_data2))
    '''
    file1 = 'model_path/pooled_feat/score_gt_%.0f_10.json' % meta_type
    data = json.load(open(file1, 'r'))
    cls_vector_1 = np.array(data['cls_vector']).reshape(-1, 2048)
    cls_gt_1 = np.array(data['cls_gt'])
    keep = cls_gt_1[:, 0] > 0
    cls_vector_1 = cls_vector_1[keep, :]
    cls_gt_1 = cls_gt_1[keep]
    
    file2 = 'model_path/pooled_feat/score_gt_%.0f_10_meta.json' % meta_type
    data = json.load(open(file2, 'r'))
    cls_vector_2 = np.array(data['cls_vector']).reshape(-1, 2048)
    cls_gt_2 = np.array(data['cls_gt'])
    keep = cls_gt_2[:, 0] > 0
    cls_vector_2 = cls_vector_2[keep, :]
    cls_gt_2 = cls_gt_2[keep]
    
    meta_len = cls_gt_2.shape[0]
    
    att_data = np.vstack((cls_vector_1, cls_vector_2))
    
    t1 = time.time()
    print('tsne start')
    tsne = TSNE(n_components = 2).fit_transform(att_data)
    t3 = time.time()
    print('tsne time used: %.2f' % float(t3 - t1))
    
    color = ['#66FFFF',
    '#808080',
    '#00FFFF',
    '#7FFFD4',
    '#008000',
    '#90EE90',
    '#DC143C',
    '#FFA500',
    '#FFEBCD',
    '#0000FF',
    '#8A2BE2',
    '#A52A2A',
    '#DEB887',
    '#5F9EA0',
    '#FF7F50',
    '#D2691E',
    '#7FFF00',
    '#6495ED',
    '#000000',
    '#D3D3D3',
    '#FFD700']

    plt.cla()
    
    tsne_1 = tsne[:-meta_len, :]
    tsne_2 = tsne[-meta_len:, :]
    '''
    for cls in range(20):
        keep = cls_gt_all[:, 0] == cls

        x = tsne[keep, 0]
        y = tsne[keep, 1]
        x = tsne[cls * shots : (cls + 1) * shots, 0]
        y = tsne[cls * shots : (cls + 1) * shots, 1]
            
        plt.scatter(x, y, s=10, c = color[cls + 1], marker = 'x')
        
    plt.legend(['aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog', 'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow', 'motorbike', 'sofa'], bbox_to_anchor=(0.95, 1))
    plt.savefig('attentions/attention_tsne_result_%.0f_epoch_%.0f.jpg' % (shots, epoch), dpi=1000)
    
    for cls in range(20):
        x = tsne[200 + cls * shots_2 : 200 + (cls + 1) * shots_2, 0]
        y = tsne[200 + cls * shots_2 : 200 + (cls + 1) * shots_2, 1]
            
        plt.scatter(x, y, s=1, c = color[cls + 1], alpha=0.5)
    '''
    for cls in range(1, 21):
        keep = cls_gt_2[:, 0] == cls

        x = tsne_2[keep, 0]
        y = tsne_2[keep, 1]
            
        plt.scatter(x, y, s=10, c = color[cls], marker = 'x', alpha=0.1)
        
    if meta_type == 1:
        plt.legend(['aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog', 'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow', 'motorbike', 'sofa'], bbox_to_anchor=(0.95, 1))
    elif meta_type == 2:
        plt.legend(['bicycle', 'bird', 'boat', 'bus', 'car', 'cat', 'chair', 'diningtable', 'dog', 'motorbike',
                             'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'aeroplane', 'bottle', 'cow',
                             'horse', 'sofa'], bbox_to_anchor=(0.95, 1))
    elif meta_type == 3:
        plt.legend(['aeroplane', 'bicycle', 'bird', 'bottle', 'bus', 'car', 'chair', 'cow', 'diningtable',
                             'dog', 'horse', 'person', 'pottedplant', 'train', 'tvmonitor', 'boat', 'cat', 'motorbike',
                             'sheep', 'sofa'], bbox_to_anchor=(0.95, 1))
    plt.savefig('attentions/meta_tsne_result_%.0f.jpg' % meta_type, dpi=1000)
    
    for cls in range(1, 21):
        keep = cls_gt_1[:, 0] == cls

        x = tsne_1[keep, 0]
        y = tsne_1[keep, 1]
            
        plt.scatter(x, y, s=1, c = color[cls], alpha=0.5)
    
    plt.savefig('attentions/meta_tsne_result_combine_%.0f.jpg' % meta_type, dpi=1000)
    t4 = time.time()
    print('scatter plot time used: %.2f' % (t4 - t3))